#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2, PointField
import numpy as np
import struct
import open3d as o3d
import time
import pandas as pd
import matplotlib.cm as cm

# --- 函数定义部分 (保持不变) ---
def pointcloud2_to_numpy(msg: PointCloud2) -> np.ndarray:
    # ... (代码与之前版本相同，此处省略以保持简洁)
    num_points = msg.width * msg.height; point_step = msg.point_step; data = msg.data; offsets = {}; dtypes = {PointField.INT8: 'b', PointField.UINT8: 'B', PointField.INT16: 'h', PointField.UINT16: 'H', PointField.INT32: 'i', PointField.UINT32: 'I', PointField.FLOAT32: 'f', PointField.FLOAT64: 'd'};
    for field in msg.fields:
        if field.name in ['x', 'y', 'z', 'intensity', 'reflectivity']:
            offsets[field.name] = field.offset; dtypes[field.name] = dtypes[field.datatype]
    if not all(k in offsets for k in ['x', 'y', 'z']): raise ValueError("PointCloud2消息缺少x, y, 或 z字段")
    has_intensity = 'intensity' in offsets or 'reflectivity' in offsets; intensity_key = 'intensity' if 'intensity' in offsets else 'reflectivity'
    points = [];
    for i in range(num_points):
        base_idx = i * point_step; x = struct.unpack_from(dtypes['x'], data, offset=base_idx + offsets['x'])[0]; y = struct.unpack_from(dtypes['y'], data, offset=base_idx + offsets['y'])[0]; z = struct.unpack_from(dtypes['z'], data, offset=base_idx + offsets['z'])[0]; intensity = 0.0
        if has_intensity: intensity = struct.unpack_from(dtypes[intensity_key], data, offset=base_idx + offsets[intensity_key])[0]
        points.append([x, y, z, intensity])
    return np.array(points, dtype=np.float32)

FIELDS_OUT = [PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1), PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1), PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1), PointField(name='intensity', offset=12, datatype=PointField.FLOAT32, count=1)]
POINT_STEP_OUT = 16

def create_pointcloud2_msg(header, points_array):
    # ... (代码与之前版本相同，此处省略以保持简洁)
    pcl_msg = PointCloud2(); pcl_msg.header = header; pcl_msg.height = 1; pcl_msg.width = len(points_array); pcl_msg.fields = FIELDS_OUT; pcl_msg.is_bigendian = False; pcl_msg.point_step = POINT_STEP_OUT; pcl_msg.row_step = POINT_STEP_OUT * pcl_msg.width; pcl_msg.is_dense = False; buffer = bytearray(pcl_msg.row_step);
    for i, point in enumerate(points_array): struct.pack_into('ffff', buffer, i * POINT_STEP_OUT, point[0], point[1], point[2], point[3])
    pcl_msg.data = buffer
    return pcl_msg

class LidarProcessor(Node):
    def __init__(self):
        super().__init__('lidar_processor')
        # ... ROS与可视化初始化 ...
        self.declare_parameter('input_cloud', '/livox/lidar')
        input_topic = self.get_parameter('input_cloud').get_parameter_value().string_value
        self.sub = self.create_subscription(PointCloud2, input_topic, self.callback, 10)
        self.pcl_pub = self.create_publisher(PointCloud2, '/livox/pointcloud2_processed', 10)
        self.elevation_map_pub = self.create_publisher(PointCloud2, '/elevation_map', 10)
        self.get_logger().info(f"订阅: {input_topic}，发布可视化点云到: /livox/pointcloud2_processed，发布高程地图到: /elevation_map")
        self.pcd_main = o3d.geometry.PointCloud(); self.vis_main = o3d.visualization.Visualizer(); self.vis_main.create_window(window_name='Main View')
        opt_main = self.vis_main.get_render_option(); opt_main.point_size = 3.0; opt_main.background_color = np.array([0, 0, 0])
        self.is_first_frame = True
        self.pcd_map = o3d.geometry.PointCloud(); self.vis_map = o3d.visualization.Visualizer(); self.vis_map.create_window(window_name='Elevation Map View')
        opt_map = self.vis_map.get_render_option(); opt_map.point_size = 5.0; opt_map.background_color = np.array([0.1, 0.1, 0.1])
        self.is_first_map_frame = True
        
        # --- 可调参数 ---
        self.frame_count = 0; self.estimated_ground_x = None; self.estimated_hook_center = None
        self.enable_ground_removal = True; self.ground_search_radius = 5.0; self.ground_ransac_distance = 0.2; self.ground_min_points_ratio = 0.05; self.ground_level_smoothing_factor = 0.1
        self.coarse_search_radius = 1.0; self.coarse_search_min_x = 0.5; self.coarse_search_max_x = 50.0
        self.rope_plane_ransac_distance = 0.1; self.rope_min_points = 30
        self.hook_position_smoothing_factor = 0.2
        self.hook_refinement_radius = 0.6; self.hook_refinement_cluster_eps = 0.2; self.hook_refinement_min_points = 15
        
        # 【修改】移除了 rope_sor 相关的参数

        self.enable_tilt_detection = True
        self.enable_elevation_map = True; self.enable_map_visualization = True
        self.map_resolution = 0.2; self.map_size_meters = 50.0; self.map_height_method = 'max'
        self.map_visualization_style = 'falt'; self.map_flat_colormap = 'terrain'; self.map_pillar_samples = 10
        self.map_num_cells = int(self.map_size_meters / self.map_resolution)

        # 【新增】高程地图中心屏蔽参数
        self.map_center_blinding_radius = 1.0  # (单位：米) 在此半径内的区域将被强制设为地面高度

        self.vis_main_view = {"front": [0.0, 0.0, -1.0], "lookat": [10.0, 0.0, 0.0], "up": [0.0, -1.0, 0.0], "zoom": 0.2}

    def callback(self, msg: PointCloud2):
        # ... (回调函数的主体处理逻辑) ...
        start_time = time.time();
        try: points_array = pointcloud2_to_numpy(msg)
        except (ValueError, KeyError) as e: self.get_logger().error(f"解析PointCloud2消息失败: {e}"); return
        if points_array.shape[0] == 0: self.get_logger().warn("接收到空点云，跳过处理。"); return
        num_original_points = len(points_array); point_labels = np.zeros(num_original_points, dtype=int); non_ground_mask = np.ones(num_original_points, dtype=bool)
        if self.enable_ground_removal:
            # ... (地面分割逻辑不变) ...
            dist_yz_all = np.sqrt(points_array[:, 1]**2 + points_array[:, 2]**2); ground_search_mask = (dist_yz_all <= self.ground_search_radius); points_for_ground_search = points_array[ground_search_mask]
            if len(points_for_ground_search) > 100:
                pcd_ground_search = o3d.geometry.PointCloud(); pcd_ground_search.points = o3d.utility.Vector3dVector(points_for_ground_search[:, :3]); plane_model, inliers = pcd_ground_search.segment_plane(distance_threshold=self.ground_ransac_distance, ransac_n=3, num_iterations=100)
                normal_vector = plane_model[:3]; is_horizontal_plane = np.abs(normal_vector[0]) > 0.9
                if len(inliers) > len(points_for_ground_search) * self.ground_min_points_ratio and is_horizontal_plane:
                    original_indices_of_ground_search = np.where(ground_search_mask)[0]; ground_indices = original_indices_of_ground_search[inliers]; point_labels[ground_indices] = 1; non_ground_mask[ground_indices] = False
                    current_ground_x = np.median(points_array[ground_indices, 0]);
                    if self.estimated_ground_x is None: self.estimated_ground_x = current_ground_x
                    else: self.estimated_ground_x = self.ground_level_smoothing_factor * current_ground_x + (1 - self.ground_level_smoothing_factor) * self.estimated_ground_x
        non_ground_points = points_array[non_ground_mask]; original_indices_of_non_ground = np.where(non_ground_mask)[0]
        dist_yz_non_ground = np.sqrt(non_ground_points[:, 1]**2 + non_ground_points[:, 2]**2); coarse_roi_mask_local = (dist_yz_non_ground <= self.coarse_search_radius) & (non_ground_points[:, 0] >= self.coarse_search_min_x) & (non_ground_points[:, 0] <= self.coarse_search_max_x)
        points_in_coarse_roi = non_ground_points[coarse_roi_mask_local]; coarse_roi_indices_global = original_indices_of_non_ground[coarse_roi_mask_local]; point_labels[coarse_roi_indices_global] = 4
        if len(points_in_coarse_roi) > self.rope_min_points:
            pcd_coarse_roi = o3d.geometry.PointCloud(); pcd_coarse_roi.points = o3d.utility.Vector3dVector(points_in_coarse_roi[:, :3]); plane_model, inliers = pcd_coarse_roi.segment_plane(distance_threshold=self.rope_plane_ransac_distance, ransac_n=3, num_iterations=100)
            normal_vector = plane_model[:3]; is_vertical_plane = np.abs(normal_vector[0]) < 0.1
            if len(inliers) > self.rope_min_points and is_vertical_plane:
                # 【修改】移除了SOR滤波，直接使用RANSAC的结果
                rope_plane_indices_in_coarse_roi = inliers
                rope_plane_indices_local = np.where(coarse_roi_mask_local)[0][rope_plane_indices_in_coarse_roi]; rope_plane_indices_global = original_indices_of_non_ground[rope_plane_indices_local]; point_labels[rope_plane_indices_global] = 3
                rope_plane_points = non_ground_points[rope_plane_indices_local]
                if self.enable_tilt_detection: self._calculate_rope_tilt(rope_plane_points)
                max_x_rope = np.max(rope_plane_points[:, 0]); coarse_hook_points = rope_plane_points[rope_plane_points[:, 0] >= (max_x_rope - 0.3)]
                if len(coarse_hook_points) > 0:
                    current_hook_center = np.mean(coarse_hook_points[:, :3], axis=0)
                    if self.estimated_hook_center is None: self.estimated_hook_center = current_hook_center
                    else: self.estimated_hook_center = self.hook_position_smoothing_factor * current_hook_center + (1 - self.hook_position_smoothing_factor) * self.estimated_hook_center
                    distances_to_stable_center = np.linalg.norm(non_ground_points[:, :3] - self.estimated_hook_center, axis=1); refinement_mask_local = (distances_to_stable_center <= self.hook_refinement_radius)
                    points_for_refinement = non_ground_points[refinement_mask_local]
                    if len(points_for_refinement) > self.hook_refinement_min_points:
                        pcd_refinement = o3d.geometry.PointCloud(); pcd_refinement.points = o3d.utility.Vector3dVector(points_for_refinement[:, :3]); cluster_labels = np.array(pcd_refinement.cluster_dbscan(eps=self.hook_refinement_cluster_eps, min_points=self.hook_refinement_min_points, print_progress=False))
                        if cluster_labels.max() > -1:
                            unique_labels, counts = np.unique(cluster_labels[cluster_labels > -1], return_counts=True); final_hook_label = unique_labels[counts.argmax()]
                            final_hook_mask_in_refinement = (cluster_labels == final_hook_label); original_indices_of_refinement = np.where(refinement_mask_local)[0]; final_hook_indices_local = original_indices_of_refinement[final_hook_mask_in_refinement]
                            final_hook_indices_global = original_indices_of_non_ground[final_hook_indices_local]; point_labels[final_hook_indices_global] = 2
                            final_hook_depth = np.max(points_array[final_hook_indices_global, 0])
                            if self.estimated_ground_x is not None:
                                height_from_ground = self.estimated_ground_x - final_hook_depth
                                self.get_logger().info(f"--- 吊钩高度(距地面): {height_from_ground:.2f} 米 ---")
        final_colors = np.zeros((num_original_points, 3)); final_colors[point_labels == 0] = [0.25, 0.25, 0.25]; final_colors[point_labels == 1] = [0.0, 0.5, 0.0]; final_colors[point_labels == 4] = [0.5, 0.5, 0.5]; final_colors[point_labels == 3] = [0.0, 0.5, 1.0]; final_colors[point_labels == 2] = [1.0, 0.0, 0.0]
        map_pcd_array = None
        if self.enable_elevation_map: map_pcd_array = self.generate_and_publish_elevation_map(points_array, point_labels, msg.header)
        self._update_visualizers(points_array, final_colors, map_pcd_array)
        process_time = (time.time() - start_time) * 1000
        self.get_logger().info(f"帧 {self.frame_count+1}: 总点数{num_original_points} → 耗时{process_time:.2f}ms\n")
        self.frame_count += 1
    
    # ... (_calculate_rope_tilt, _update_visualizers, _colorize_flat_map, _create_pillar_map 函数与之前版本相同，省略) ...
    def _calculate_rope_tilt(self, rope_points):
        if len(rope_points) < 10: return
        centroid = np.mean(rope_points[:, :3], axis=0); centered_points = rope_points[:, :3] - centroid; cov_matrix = np.cov(centered_points, rowvar=False)
        _, _, vh = np.linalg.svd(cov_matrix); principal_direction = vh[0, :]
        vertical_axis = np.array([1.0, 0.0, 0.0]); cos_angle = np.abs(np.dot(principal_direction, vertical_axis)); cos_angle = np.clip(cos_angle, -1.0, 1.0)
        angle_rad = np.arccos(cos_angle); angle_deg = np.degrees(angle_rad)
        self.get_logger().info(f"--- 吊钩倾斜角度: {angle_deg:.2f} 度 (0度为垂直) ---")
    def _update_visualizers(self, main_points, main_colors, map_pcd_array):
        self.pcd_main.points = o3d.utility.Vector3dVector(main_points[:, :3]); self.pcd_main.colors = o3d.utility.Vector3dVector(main_colors)
        if self.is_first_frame:
            if len(self.pcd_main.points) > 0:
                self.vis_main.add_geometry(self.pcd_main)
                ctr = self.vis_main.get_view_control(); ctr.set_front(self.vis_main_view["front"]); ctr.set_lookat(self.vis_main_view["lookat"]); ctr.set_up(self.vis_main_view["up"]); ctr.set_zoom(self.vis_main_view["zoom"])
                self.is_first_frame = False
        else: self.vis_main.update_geometry(self.pcd_main)
        self.vis_main.poll_events(); self.vis_main.update_renderer()
        if self.enable_map_visualization:
            map_points_to_show, map_colors_to_show = np.array([]), np.array([])
            if map_pcd_array is not None and len(map_pcd_array) > 0:
                if self.map_visualization_style == 'pillar': map_points_to_show, map_colors_to_show = self._create_pillar_map(map_pcd_array)
                else: map_points_to_show = map_pcd_array; map_colors_to_show = self._colorize_flat_map(map_pcd_array)
            if len(map_points_to_show) > 0:
                self.pcd_map.points = o3d.utility.Vector3dVector(map_points_to_show); self.pcd_map.colors = o3d.utility.Vector3dVector(map_colors_to_show)
                if self.is_first_map_frame:
                    self.vis_map.add_geometry(self.pcd_map); ctr = self.vis_map.get_view_control(); ctr.set_front([-1, 0, 0]); ctr.set_up([0, 0, 1]); self.is_first_map_frame = False
                else: self.vis_map.update_geometry(self.pcd_map)
            else: self.pcd_map.clear(); self.vis_map.update_geometry(self.pcd_map)
            self.vis_map.poll_events(); self.vis_map.update_renderer()
    def _colorize_flat_map(self, map_pcd_array):
        heights = map_pcd_array[:, 0]; min_h, max_h = np.min(heights), np.max(heights)
        if max_h == min_h: normalized_heights = np.full_like(heights, 0.5)
        else: normalized_heights = (heights - min_h) / (max_h - min_h)
        colormap = cm.get_cmap(self.map_flat_colormap); colors = colormap(normalized_heights)[:, :3]
        return colors
    def _create_pillar_map(self, map_pcd_array):
        if self.estimated_ground_x is None: return np.array([]), np.array([])
        all_pillar_points, all_pillar_colors = [], []
        for i in range(len(map_pcd_array)):
            height, y, z = map_pcd_array[i]
            pillar_x = np.linspace(height, self.estimated_ground_x, self.map_pillar_samples); pillar_y = np.full(self.map_pillar_samples, y); pillar_z = np.full(self.map_pillar_samples, z)
            pillar_points = np.vstack([pillar_x, pillar_y, pillar_z]).T; all_pillar_points.append(pillar_points)
            pillar_colors = np.full_like(pillar_points, [0.0, 0.5, 1.0]); pillar_colors[0, :] = [1.0, 0.0, 0.0]; all_pillar_colors.append(pillar_colors)
        if not all_pillar_points: return np.array([]), np.array([])
        return np.vstack(all_pillar_points), np.vstack(all_pillar_colors)

    def generate_and_publish_elevation_map(self, points, labels, header):
        # ... (与之前版本相同，但加入了绳索点和中心屏蔽)
        map_points_mask = (labels == 1) | (labels == 0) | (labels == 4) | (labels == 3)
        map_points = points[map_points_mask]
        if len(map_points) < 10: return None
        y_coords = map_points[:, 1]; z_coords = map_points[:, 2]; x_heights = map_points[:, 0]
        center_y = np.median(y_coords); center_z = np.median(z_coords)
        half_size = self.map_size_meters / 2.0
        relative_y = y_coords - center_y; relative_z = z_coords - center_z
        col_indices = ((relative_y + half_size) / self.map_resolution).astype(int); row_indices = ((relative_z + half_size) / self.map_resolution).astype(int)
        valid_indices_mask = (col_indices >= 0) & (col_indices < self.map_num_cells) & (row_indices >= 0) & (row_indices < self.map_num_cells)
        if not np.any(valid_indices_mask): return None
        df = pd.DataFrame({'row': row_indices[valid_indices_mask], 'col': col_indices[valid_indices_mask], 'height': x_heights[valid_indices_mask]})
        if self.map_height_method == 'max': grid_heights = df.groupby(['row', 'col'])['height'].max()
        elif self.map_height_method == 'min': grid_heights = df.groupby(['row', 'col'])['height'].min()
        else: grid_heights = df.groupby(['row', 'col'])['height'].mean()
        if grid_heights.empty: return None
        
        map_grid_points = []
        for (r, c), h in grid_heights.items():
            y_rel = (c + 0.5) * self.map_resolution - half_size; z_rel = (r + 0.5) * self.map_resolution - half_size
            y_world = y_rel + center_y; z_world = z_rel + center_z
            map_grid_points.append([h, y_world, z_world])
        if not map_grid_points: return None
        map_pcd_array = np.array(map_grid_points, dtype=np.float32)

        # 【核心修改】执行中心区域屏蔽
        if self.map_center_blinding_radius > 0 and self.estimated_ground_x is not None:
            # 计算每个栅格中心到LiDAR原点(Y=0, Z=0)的水平距离
            # 注意 map_pcd_array 的列是 [X, Y, Z]
            grid_center_distances = np.sqrt(map_pcd_array[:, 1]**2 + map_pcd_array[:, 2]**2)
            blinding_mask = grid_center_distances <= self.map_center_blinding_radius
            # 将屏蔽区域内栅格的高度强制设为地面高度
            map_pcd_array[blinding_mask, 0] = self.estimated_ground_x
            self.get_logger().info(f"高程地图中心屏蔽：{np.sum(blinding_mask)}个栅格被强制设为地面高度。")

        publish_array = np.hstack([map_pcd_array, map_pcd_array[:, [0]]])
        map_msg = create_pointcloud2_msg(header, publish_array)
        self.elevation_map_pub.publish(map_msg)
        return map_pcd_array

    def destroy_node(self):
        # ... (与之前版本相同)
        self.get_logger().info("正在关闭节点..."); self.vis_main.destroy_window(); self.vis_map.destroy_window(); super().destroy_node(); self.get_logger().info("节点已关闭。")

def main(args=None):
    # ... (与之前版本相同)
    rclpy.init(args=args); node = LidarProcessor()
    try: rclpy.spin(node)
    except KeyboardInterrupt: node.get_logger().info("捕获到Ctrl+C，开始关闭...")
    finally: node.destroy_node(); rclpy.shutdown()

if __name__ == '__main__':
    main()